from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers
from gmutil import SM2PrivateKey,SM2PublicKey
from cryptography.hazmat.primitives import serialization
from os.path import join, dirname, realpath
from datetime import datetime, timedelta
import base64
import requests
import json

cwd = dirname(realpath(__file__))

UNIVERSITY_ID_UTO = 'FF11010003'
UNIVERSITY_ID_THU = '4111010003'

SAMPLE_SERVER = r'http://main.chenqiang.me:5000'
# SAMPLE_SERVER = r'http://127.0.0.1:5000'

with open(join(dirname(realpath(__file__)), 'directory.json'), 'r', encoding='utf-8') as f:
    DIRECTORY = json.load(f)

with open(join(dirname(realpath(__file__)), 'UTO_sm2_pri.pem'), 'rb') as f:
    UTO_SM2_PRIVATE_KEY = SM2PrivateKey.from_pem(f.read())

with open(join(dirname(realpath(__file__)), 'UTO_rsa_pri.pem'), 'rb') as f:
    UTO_RSA_PRIVATE_KEY = serialization.load_pem_private_key(f.read(), None)


def request_qrcode(userid: str):
    resp = requests.get('{}/qrcode/{}'.format(SAMPLE_SERVER, userid))
    if resp.status_code == 200:
        return resp.text
    else:
        return None

def authenticate(code: str):
    acquirer_code = UNIVERSITY_ID_UTO
    barcode = code
    now = datetime.now()
    timestamp = now.strftime("%Y%m%d%H%M%S")
    to_verify = acquirer_code.encode() + barcode.encode() + timestamp.encode()
    print('请求待签名数据：', to_verify)
    sig = UTO_SM2_PRIVATE_KEY.sign(to_verify)
    print('签名值：', base64.urlsafe_b64encode(sig).decode())

    req_data ={
        'acquirer_code': acquirer_code,
        'barcode': barcode,
        'timestamp': timestamp,
        'kid': '01',
        'signature': base64.urlsafe_b64encode(sig).decode().rstrip("=")
    }

    req_headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
    }

    for algorithm in ('SM2', 'RSA'):
        req_data['algorithm'] = algorithm
        req = requests.Request('POST', f'{SAMPLE_SERVER}/auth', data=req_data, headers=req_headers)
        prepared = req.prepare()
        print('request method:', prepared.method)
        print('request url:', prepared.url)
        print('request header:', prepared.headers)
        print('request body:', prepared.body)

        s = requests.Session()
        resp = s.send(prepared)
        # resp = requests.post(f'{SAMPLE_SERVER}/auth', data=req_data, headers=req_headers)
        if resp.status_code == 200:
            resp_data = json.loads(resp.text)
            if (code := resp_data['code']) != 0:
                print(f'认证错误：{resp_data["message"]}（{code}）')
                continue

            resp_kid = resp_data.pop('kid')
            resp_sig = base64.urlsafe_b64decode(resp_data.pop('signature') + '==')

            entry = next((e for e in DIRECTORY if e['code'] == UNIVERSITY_ID_THU), None)
            if entry is None:
                return None
            key = next((k for k in entry['public_keys']['keys'] if k['kid'] == resp_kid), None)
            if key is None:
                return None

            to_verify = '&'.join(
                ['{}={}'.format(k, v) for k, v in sorted(resp_data.items(), key=lambda item: item[0])])
            print(to_verify)
            message = to_verify.encode('gb18030')
            print(resp_sig.hex())

            if algorithm == 'SGD_SM3_SM2':
                x_octets, y_octets = base64.urlsafe_b64decode(key['x']), base64.urlsafe_b64decode(key['y'])
                resp_pub_key = SM2PublicKey.from_coordinates(
                    int.from_bytes(x_octets, byteorder='big', signed=False),
                    int.from_bytes(y_octets, byteorder='big', signed=False))
                if not resp_pub_key.verify(message, resp_sig):
                    print("认证无效：SM2验签失败")
            elif algorithm == 'RS256':
                n = int.from_bytes(base64.urlsafe_b64decode(key['n']), byteorder='big', signed=False)
                e = int.from_bytes(base64.urlsafe_b64decode(key['e']), byteorder='big', signed=False)
                pubkey = RSAPublicNumbers(e, n).public_key()
                try:
                    pubkey.verify(resp_sig, message, padding.PKCS1v15(), hashes.SHA256())
                except InvalidSignature:
                    print("认证无效：RSA验签失败")
            print(json.dumps(resp_data, ensure_ascii=False, indent=4))
        else:
            print(resp.status_code)
            return None


if __name__ == '__main__':
    qrcode = request_qrcode('2007980012')
    print(qrcode)
    assert qrcode.startswith('CUCC:')
    qrdata = base64.urlsafe_b64decode(qrcode[5:])
    version = qrdata[0]
    if version != 1:
        raise ValueError('CUCC版本号错误')
    university_id = qrdata[1:6]
    print(f'发卡高校：{university_id.hex()}')
    issuing_time = qrdata[6:13]
    print(f'生成时间：{issuing_time.hex()}')

    custom_len = qrdata[13]
    if len(qrdata) != custom_len + 14:
        raise ValueError('赋码高校自定义内容长度错误')

    authenticate(qrcode[5:])
